#提取深度特征，使用机器学习模型分类


import os
from sklearn.metrics import confusion_matrix
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import numpy as np
from train_model import NewDataset,cal_patient_acc
from train_kfold import TrianDataset
from torch.utils.data import DataLoader, WeightedRandomSampler
from torch.utils.tensorboard import SummaryWriter
from patient_information import find_patient_files,load_patient_data,get_grade_number,get_murmur,get_num_locations,get_locations



if __name__ == "__main__":
    fold_path = 'data_5fold_new/3_fold'
    best_model_dir = os.path.join('SKCNN_result',fold_path,'model/best_model')
    # best_model_dir = 'ViT_result/data_5fold_new2/4_fold/model/best_model'
    feature_model_dir = 'Deep_Feature_Model'
    model = torch.load(best_model_dir)
    #将训练好的模型全连接层去掉
    # model.fc = nn.Identity()
    # model.soft = nn.Identity()
    # model.head = nn.Identity()
    # model.lin = nn.Identity()
    model.classfier = nn.Identity()
    # torch.save(model, os.path.join( feature_model_dir,'feature_extractor_vit'))

    # fold_path = 'data_5fold_new2/4_fold'
    data_path = os.path.join(fold_path, 'feat')
    label_path = os.path.join(fold_path, 'label')

    data_train = np.load(data_path + '/train_loggamma.npy', allow_pickle=True)
    data_vali = np.load(data_path + '/test_loggamma.npy', allow_pickle=True)

    label_train = np.load(label_path + '/train_label.npy', allow_pickle=True)
    label_vali = np.load(label_path + '/test_label.npy', allow_pickle=True)

    train_set = TrianDataset(wav_label=label_train, wav_data=data_train)
    vali_set = TrianDataset(wav_label=label_vali, wav_data=data_vali)

    batch_size = 1
    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, drop_last=True)
    test_loader = DataLoader(vali_set, batch_size=batch_size, shuffle=True, drop_last=True)

    # data_train = (torch.from_numpy(data_train)).float()
    # input = torch.rand(128, 32, 239)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    model.eval()

    train_feature = []
    vali_feature = []

    all_y_pred = []  # 存放3s样本的预测概率
    all_y_pred_label = []  # 存放3s样本的真实标签
    all_label = []
    with torch.no_grad():
        for i, data in enumerate(train_loader):
            x, y = data
            x = x.to(device)
            feature = model(x)
            # _, y_pred = outputs.max(1)
            # all_label.append(y.cpu().detach())
            # all_y_pred_label.append(y_pred.cpu().detach())

            train_feature.append(feature.cpu().detach())
    train_feature = np.vstack(train_feature)

    # y_pred = np.array(np.hstack(all_y_pred_label))
    # y_true = np.array(np.hstack(all_label))
    # # 计算混淆矩阵
    # cm = confusion_matrix(y_true, y_pred)
    np.save(data_path + r'/train_SKCNN_feature.npy', train_feature)

    all_y_pred = []  # 存放3s样本的预测概率
    all_y_pred_label = []  # 存放3s样本的真实标签
    all_label = []
    with torch.no_grad():
        for i, data in enumerate(test_loader):
            x, y = data
            x = x.to(device)
            feature = model(x)
            # _, y_pred = outputs.max(1)
            # all_label.append(y.cpu().detach())
            # all_y_pred_label.append(y_pred.cpu().detach())

            vali_feature.append(feature.cpu().detach())

    vali_feature = np.vstack(vali_feature)

    # y_pred = np.array(np.hstack(all_y_pred_label))
    # y_true = np.array(np.hstack(all_label))
    # # 计算混淆矩阵
    # cm2 = confusion_matrix(y_true, y_pred)
    np.save(data_path + r'/vali_SKCNN_feature.npy', vali_feature)


    print(train_feature.shape)
